AsyncGRPO: train OpenEnv harnesses from validated token captures - #6947
adithya-s-k wants to merge 51 commits into
Conversation
…hrough OpenEnv Trains against a Harbor task through mini-swe-agent running in an E2B sandbox. The agent owns its own loop; TRL stands up an endpoint, lets it drive, and reads back the captured token ids and logprobs. That is what makes an installed harness trainable without reimplementing it, and it is the difference from examples/grpo_harbor, which runs Harbor tasks against harnesses written inside TRL with TRL owning the loop. mini-swe-agent is the default on measured grounds rather than taste: across a 15-harness sweep on the same 50 tasks it was the most accurate and the most turn-efficient, its prompt re-render is byte-exact against the engine's prompt_token_ids, and it is the only harness that can express a step limit. The re-render matters because TRL rebuilds each prompt locally, and for three of twelve harnesses measured that drifts (claude-code +2 tokens, gemini-cli +2, kimi-cli -10 per tool call) -- invisible for eval, forking the trajectory every turn when training. The step limit is not a cost control. Every turn re-sends the whole conversation, so a rollout's packed length grows with the SQUARE of its turn count; unbounded 58-turn rollouts were enough to OOM the loss step on an 80 GiB card. The docstring states one caveat rather than hiding it: on this path HarnessRolloutOutcome carries a single verifier scalar, not the component dict, so a 'submission' term giving partial credit is unavailable and the reward is all-or-nothing. On a suite the model solves ~16% of the time that means most groups score identically and those steps teach nothing, so the example says to shape component rewards where the suite emits them and otherwise to pick tasks the model solves sometimes.
…d note The grpo_harbor reference would dangle once that example is deprecated, and a docstring should not point at something being removed. The reward paragraph is cut to what it is -- correctness plus a correctness-gated efficiency term, with --reward-key for suites that emit a dict -- rather than a discussion of what the single-scalar path cannot express.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 025f3309ee
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Brings the whole stack up inside one container: openenv harbor serve on CPU, vllm serve on GPU 0, the trainer on GPU 1. Self-contained rather than importing its sibling, because hf jobs uv run uploads a single script. Everything lives in one job because AsyncGRPO syncs weights into vLLM over NCCL, which needs both on the same host's GPUs. That leaves only the OpenEnv server placeable, and keeping it here puts the proxy's hop to vLLM on localhost; hosting it on a Space via `openenv harbor push` works but adds a public hop to every model call. The tunnel is not a workaround for missing infrastructure. Jobs can publish a port at <job_id>--<port>.hf.jobs, but access needs an HF token, and the agent's Authorization header already carries its rollout session key -- that key IS how the proxy routes concurrent rollouts, so it cannot carry a second credential. An unauthenticated outbound tunnel needs no ingress at all. The readiness check verifies the tunnel SERVES THE PROXY, not that a port answers. A tunnel whose forwarding process dies keeps resolving and returns the provider's error page; agents then get HTML where an OpenAI endpoint should be, make zero model calls, and every rollout comes back unscorable while the server's own /health still reports healthy. The URL is only accepted after a request through it returns our health document. The parser was replayed over 32 real server logs and matches all 12 published URLs, gradio and cloudflare alike. The mounted bucket holds checkpoints and HF_HOME, and says so loudly when absent rather than losing them silently. Deliberately not the sandbox templates: those are keyed by image hash on the provider's side and already persist across jobs, which is the expensive warm step.
…itten
--split was required, so the reproduction command carried a placeholder nobody could paste. It now
defaults to a public Harbor suite (AdithyaSK/data_agent_rl_environment_train, verified public), which
means the whole thing is one copy-pasteable line with no <angle brackets> in it:
hf jobs uv run --flavor h200x2 --image huggingface/trl \
--secrets HF_TOKEN --secrets E2B_API_KEY \
https://raw.githubusercontent.com/.../async_grpo_harbor_hf_jobs.py
The docstring shows that form first and the bucket form second, since the bucket is optional -- without
it the job still trains, it just loses its checkpoints when the container goes away, which the script
already warns about at startup.
Also fixes a flag that could not do anything: --no-enable-thinking was declared with store_false over a
default of False, so it could only set what was already set. Thinking is off by default because the
trainer's chat_template_kwargs must agree with how the engine was served, and --enable-thinking is now
the switch that changes it.
Found by auditing the script against the Jobs environment rather than the cluster it was written on.
1. vllm was never declared. `uv run --script` builds an isolated environment, so the base image's vllm
is not guaranteed importable -- and the script's own PATH guard would then have blamed "the PEP 723
dependencies did not install", which would have been true but unhelpful.
2. No sandbox template warm. This is the one that would have looked like a harness bug. Harbor decides
whether to build from alias_exists(), which flips true when a build STARTS, so num_generations
rollouts racing their first visit to a task all see "exists" and fail against a half-built image with
404: tag 'default' does not exist. On a cold template with 8 concurrent generations that is the
common case. One serial `openenv harbor rollout` first makes the build serial and every later rollout
finds a finished image; the provider keys images by content hash, so later jobs pay nothing.
Non-fatal on failure: the warm rollout can fail for reasons that say nothing about training, and the
build it triggered still happened.
3. No GPU-count guard. On a one-GPU flavor the trainer was pointed at device 1 and would have failed
deep inside CUDA instead of at startup with a sentence naming --flavor h200x2.
4. The run name fell back to a constant ("local") when HF_JOB_ID was absent -- and that variable's name
is not something this script can verify. Trackio keys a run by name inside a project, so the constant
would fold separate jobs into one history, worst exactly when relaunching after a failure. It now
falls back to a timestamp.
sergiopaniego
left a comment
There was a problem hiding this comment.
quick review, first pass, we need to add the example to docs/source/example_overview.md
Restructured to match the pattern the opencode recipe already uses for Hugging Face Jobs (described in sergiopaniego's TRL x OpenEnv x Harbor post, with the launcher published as a gist): a small launcher wraps the processes a Job cannot start on its own, and the training script is DOWNLOADED rather than duplicated. That last part is the reason for the change. The previous commit added a second 489-line file that restated the whole training path, so the two could drift and a reader had to diff them to see what was actually different. Now async_grpo_harbor.py is the single canonical script, byte-identical whether it runs locally or in a Job, and launcher.py holds only what is genuinely Jobs-specific. Ours needs three processes where the opencode launcher needs two, because the Harbor dataset and the capture proxy are served by openenv harbor serve. It also tunnels a different hop: opencode tunnels vLLM, since its in-sandbox proxy calls the engine directly, whereas here the sandboxed agent calls the capture proxy and the engine stays entirely private on localhost. --train-script-url exists because the equivalent opencode launcher still points at examples/scripts/openenv/opencode_hf_sandbox.py, a pre-reorg path that now 404s. A stale URL surfaces as a download failure minutes into a paid job, so this one is overridable and can be pointed at a branch before the example is merged. Carried over from the audit of the deleted file: vllm is declared, the sandbox template is warmed by one serial rollout before any group runs concurrently, there is a GPU-count guard, the tunnel readiness check verifies a request THROUGH the tunnel reaches the proxy, and the run name never falls back to a constant.
Verifying once at startup is not enough, and this is measured rather than defensive. Over 27 hours on our own cluster the published tunnel stopped serving the capture proxy 69 times -- roughly once every 24 minutes -- so a run of any length loses it mid-flight. When that happens the sandboxed agent gets the tunnel provider's error page instead of an OpenAI endpoint, makes zero model calls, and every rollout comes back unscorable, with nothing in the trainer's logs to say why. A daemon thread re-probes the published URL and restarts the Harbor server after two consecutive failures, so one flaky request cannot bounce a healthy server mid-step. The restart changes the published URL, which is fine: the trainer talks to the server over localhost and the server hands its current URL to each new sandbox, so only in-flight rollouts are lost. Which check does the work is worth recording. The failure signature originally debugged -- the provider's "no interface is running" placeholder -- was 2 of those 69. The other 137 probe failures were plain 502s. So the test is positive and generic: the URL must return OUR health document. Enumerating known failure modes would have caught almost none of them.
All four test suites on huggingface#6947 fail on one assertion, and it is not a code failure: tests/test_examples_index.py::test_examples_index_matches_folders AssertionError: Example folders missing from the Index table in example_overview.md: {'async_grpo_harbor'} `test_examples_index_matches_folders` requires every directory under `examples/` to have a matching row in the index, so adding the example without the row turns all of "latest", "dev", "minimum versions" and "without optional dependencies" red at once. One row fixes all four. Placed before `async_grpo_math` to keep the table alphabetical, matching the neighbouring `async_grpo_opencode` entry's format.
|
Nice!
Also the docstring says "a 15-harness sweep" then "three of the twelve harnesses measured". which is it, and is the sweep written up anywhere we can link? Feel free to merge |
…prompt In the loop-owning path the agent runs its own tool loop in a sandbox and TRL reads back what it did. Rebuilding each turn's prompt with apply_chat_template to recover the token ids produces a DIFFERENT string from the one the engine scored. Measured on Qwen3.5-4B, the rebuilt prompt matched the engine on 0 of 28 turns. Nothing errors when that happens, which is the problem. Turns that should chain by exact token prefix instead look like unrelated short rollouts, so one conversation fragments into several and every fragment still trains. Two production runs spent a night each in exactly that state with healthy-looking reward curves; the tell was samples_per_rollout equalling turns_mean exactly, i.e. 100% forking. _turns_from_trace now takes `prompt_token_ids` from the trace entry and raises when it is absent, naming the two vLLM flags that produce it. A hard failure rather than a fallback on purpose: the fallback IS the bug, and it is invisible. An engine without --return-tokens-as-token-ids yields rollouts that look completely normal and carry nothing to train on, so the run has to stop at the first rollout instead of discovering it at the first weight update. The tokenizer argument is gone from that path -- there is nothing left for it to do. TurnRecord gains `output_mask`, and _SampleBuilder.append_turn honours it instead of hardcoding 1 across the completion. This is not cosmetic: capture masks a turn OUT when its logprobs were rejected on ingest, while keeping its tokens as context and zero-filling the logprobs. Without the mask those positions train against logprob 0.0, i.e. p=1.0 -- a confident target for a turn we explicitly could not trust. Requires the matching OpenEnv change, which puts `prompt_token_ids` and `loss_mask` on TraceEntry. Against an older OpenEnv this raises immediately and says which flags are missing, which is the intended behaviour rather than a regression. Note for anyone tempted to prototype this as a monkeypatch: it cannot work. The rollout loop runs in a multiprocessing child created with `spawn`, which re-imports every module from scratch, so a parent-side rebind of a module function is simply lost. A patch that mutates the tokenizer OBJECT survives, because that is pickled into the child; one that rebinds a function does not. That asymmetry is why the re-rendering went unnoticed for so long -- the patch logged its install line and its counters never moved.
# Conflicts: # docs/source/example_overview.md
…unables, fix the harness counts From @qgallouedec's review: * THIRTY flags down to twenty-three. The seven removed (`--optim`, `--no-bf16`, `--gradient-checkpointing` / `--no-gradient-checkpointing`, `--save-steps`, `--save-total-limit`, `--seed`) only re-exported `TrainingArguments`. An example is meant to be read and edited, not driven like a product, so those are literals in the config now with a note where the choice is not obvious (gradient checkpointing is on because rollout sequences here are long enough that activations dominate). Checkpointing defaults to off, since a short example run has nothing worth keeping. * The two reward tunables are plain module constants (`W_TOOL_EFFICIENCY = 0.3`, `TOOL_BUDGET = 15.0`) rather than `os.environ.get`, matching the opencode example -- one channel for tunables. The one remaining `os.environ.get` reads `SLURM_JOB_ID`, which is a fact about the environment rather than a knob. * The harness counts were not contradictory, just under-specified: FIFTEEN harnesses were probed, TWELVE produced usable rollouts, and it is three of those twelve that had to re-render. Said explicitly now. Also resolves the `docs/source/example_overview.md` conflict with upstream: upstream added an `async_grpo_math` row while this branch added `async_grpo_harbor`. Both are kept, with upstream's version of the shared row. `tests/test_examples_index.py` passes -- that gate fails all four suites at once when a row is missing, which is what @sergiopaniego's review was about.
CI's code-quality job is `pre-commit run --all-files`, and its `doc-builder-style` hook rewraps docstrings to max_len 119. My `_turns_from_trace` docstring wrapped at a narrower width, so the hook modified a file and the job failed. Worth recording for next time: running ruff locally does NOT predict this job. ruff passed on all five of my changed files while the real failure was doc-builder rewrapping one line. The gate is pre-commit, so reproduce it with pre-commit (or at least the pinned doc-builder rev) rather than with ruff alone. The notebook errors ruff reported locally were a red herring -- those files are byte-identical to upstream and I touched no notebook.
…_fn shapes it `rollout_reward_fn` replaces `env_reward` wholesale, and the column it feeds is already named `rewards/harness_reward`. So the moment a caller passes a shaping term, `reward` and `rewards/harness_reward` are the SAME shaped number and the raw verifier score is logged nowhere. That makes every shaping term look exactly like progress. A reward curve that climbs because the policy got faster is indistinguishable from one that climbs because it solved more -- which is the one question the run exists to answer. It also silently breaks comparison against every earlier run, because the headline number changed units without changing name. Records `rollout/correctness_mean` from the same `_rates` channel `rollout/turns_mean` already uses, so it needs no new plumbing and drains generically. Only when a `rollout_reward_fn` is actually set: with no shaping term `reward` IS correctness and a second identical series would just be noise.
|
Holding off on merging this, one thing I'd like to settle first.
Also, since the producer isn't written yet, merging this first sets that contract "by accident" (rather than by agreement). I'd wait until what defines No rush from my side |
There was a problem hiding this comment.
We should update the Ready-to-use OpenEnv examples note and include more details the harbor integration, no? since the other integration will be removed in #6948
…act-v2 # Conflicts: # trl/generation/vllm_client.py
| raise CaptureContractError("loop-owning training requires OpenEnv fetch_training_trace()") | ||
| loop_session = cast(TrainableSession, session) | ||
| try: | ||
| loop_session.wait_for_completion() |
There was a problem hiding this comment.
Notblocking but just to not : Longer term, an explicit execution status / exception would be nice to make this independent of the session implementation. Having explicit status is nice, because it also means we can switch on that to know if there is a partial caputre we can train on for example ??
| trace = entries | ||
| except (ValueError, TypeError, KeyError) as exc: | ||
| raise CaptureContractError(str(exc)) from exc | ||
| completion = _messages_from_trace(entries) |
There was a problem hiding this comment.
I like that we removed the _agent_turn_fn and the other hooks, but as a followup :
OpenEnv would provide the full final assistant message directly. _messages_from_trace() assumes the last request contains the conversation. It also means TRL assumes the structure of the messages without explicit contract.
| completion = _messages_from_trace(entries) | ||
| tool_calls_by_name = _tool_call_counts_by_name(entries) | ||
| tool_call_count = sum(tool_calls_by_name.values()) | ||
| tool_failure_count = _tool_failure_count(entries) |
There was a problem hiding this comment.
This is really brittle. Same as above, we parse the entries to check if there is an error in the message. Not great ! tool outcomes should come from structured execution
records where available, as in the white-box path.
| if harness_adapter is None: | ||
| sampling_kwargs = { | ||
| "temperature": self.temperature, | ||
| "top_p": self.top_p, | ||
| "top_k": self.top_k, | ||
| "min_p": self.min_p, | ||
| "repetition_penalty": self.repetition_penalty, | ||
| } | ||
| sampling_kwargs = {key: value for key, value in sampling_kwargs.items() if value is not None} | ||
| self._factory = harness_session_factory(sampling=training_sampling(sampling_kwargs)) |
There was a problem hiding this comment.
Ok thats nice for sampling mismatch. But as a follow-up, can we make this constructor requirement explicit in the public type/ API and include execution limits?
Also, the loop-owning path currently receives sampling params but not the worker’s max_tokens nor turn limit.
| turns = _turns_from_training_trace(captured) | ||
| entries = captured.to_trace_entries() | ||
| trace = entries | ||
| except (ValueError, TypeError, KeyError) as exc: |
There was a problem hiding this comment.
another follow-up: OpenEnv should expose session-boundary error types so TRL can distinguish invalid captures from execution/infrastructure failures without interpreting generic Python exceptions.
There was a problem hiding this comment.
Thanks @adithya-s-k 🤗 .
I’m OK with the direction of this PR it cleans up a lot of stuff (mainly on the tokens/masking).
I think the remaining API improvements can be handled in a follow-up PR. Main misssing spots IMHO are session execution status and better types/contract:
-
Execution and capture lifecycle: TRL ignores the return value of
wait_for_completion()and only recognizes timeout through TimeoutError. OpenEnv should report an explicit status (completed, timed_out, budget_exhausted, cancelled, failed) and guarantee that execution has stopped and the capture is finalized before returning it. A valid partial capture should remain distinguishable from a successfully completed rollout. Then TRL can decide whether to train on it or not. Cancellation and repeated cleanup should have documented behavior, like is it thread safe cancellation ? is it async safe ? etc... -
Conversation and tool outcomes: TRL still reconstructs the conversation from the last model request, guesses tool failures from words in tool output. OpenEnv should provide full assistant messages and structured tool outcomes where available ....
-
Configuration and typing: The
sampling=constructor requirement is not clear. We have aCallable[..., ResourceSessionFactory]. A public typed configuration boundary should define sampling, execution budgets etc....
A future session result could consolidate **the status, the finalized training capture, the conversation, the tool outcomes, and the verification data((. That would let TRL focus on tensor (ids/masks), packing, and the correct loss compute which is the job of the lib IMO.
Also I guess we shouldn't merge this PR until : huggingface/OpenEnv#1280 lands ?
qgallouedec
left a comment
There was a problem hiding this comment.
Thanks, this answers my questions from the 14th: the producer owns the masks via TrainingTrace, and partial masks are supported.
Before merging:
- huggingface/OpenEnv#1280 is still open, and the examples and CI pin its commit from your fork. Once #1280 is merged, could you pin to that upstream commit (or a release) from
huggingface/OpenEnv? clear_cache=Trueinpause()is a no-op: vLLM's/pausealready defaults to it (and marks it deprecated). The docstring fix is fine, but i'd drop the param.- Could you drop
tests/test_vllm_control_requests.py? It mostly tests code this PR doesn't touch (reset_prefix_cache,_post,pause()raising on non-200 are all on main already), and it's in coretests/while targeting experimental clients. - Same for
tests/experimental/test_harbor_example.pyand the newopenenv_contractCI job: we don't test examples, only what's packaged intrl.
Then good to go on my side.
| openenv_contract: | ||
| permissions: | ||
| contents: read | ||
| name: OpenEnv capture contract (CPU) | ||
| runs-on: ubuntu-latest | ||
| env: | ||
| PYTHONPATH: temp/openenv/envs | ||
| CUDA_VISIBLE_DEVICES: "" | ||
| steps: | ||
| - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 | ||
| - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 | ||
| with: | ||
| repository: adithya-s-k/OpenEnv | ||
| ref: 58e48a95a34c0d279cfcb01c95b8688de23c79f1 | ||
| path: temp/openenv | ||
| - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 | ||
| with: | ||
| python-version: '3.12' | ||
| - uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 | ||
| - name: Install CPU test dependencies | ||
| run: | | ||
| uv venv | ||
| uv pip install torch --index-url https://download.pytorch.org/whl/cpu | ||
| uv pip install ".[test]" ./temp/openenv | ||
| - name: Check pinned producer and consumer | ||
| run: | | ||
| source .venv/bin/activate | ||
| python -c "import harbor_env.harness; import trl.experimental.async_grpo.openenv_harness" | ||
| python -m pytest -q tests/experimental/test_openenv_tito.py tests/experimental/test_harbor_example.py \ | ||
| tests/experimental/test_async_grpo_trainer.py::TestReconciler tests/test_vllm_control_requests.py | ||
|
|
There was a problem hiding this comment.
It tests an unmerged fork and example code, and it runs on CPU, so none of it fits our CI.
The one useful part is checking the trace format. That can go in the regular experimental tests once OpenEnv ships a release.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 2de2cea. Configure here.

What does this PR do?
Add a Harbor/OpenEnv example for agents that own their tool loop. OpenEnv supplies a validated
TrainingTracewith engine tokens, behavior logprobs, selected agent calls and required full-sequence masks. TRL maps the completion mask toTurnRecord.output_mask, preserving partial masks and shared-prefix supervision without re-tokenizing or selecting turns through hooks.Examples now pin the merged upstream OpenEnv commit
86a180ede21e044f7929b9a7783ad83aa67d83a3from OpenEnv #1280. The factory receives the worker's sampling policy before generation;verify()remains the correctness source. The Harbor example includes local and HF Jobs setup, checkpoint saving and metadata-session cleanup. Native OpenCode uses the same typed contract and performs CLI warmup in the parent before spawning workers.Quentin's requested example tests and dedicated CPU CI job are removed. Packaged contract tests remain under
tests/experimental, with optional producer imports skipped when unavailable. Both experimental vLLM clients retainmode=keepand omit the redundantclear_cacheargument.Amine's execution status/finalization, structured conversations/tool outcomes, typed execution limits and session error types remain follow-up API work. The current timeout and budget behavior is documented. Rollout-normalized weighting and tree attention are also outside this PR.
Validation: merged current TRL main. Against the merged OpenEnv producer, 51 packaged contract/reconciliation tests and 26 local launcher/HTTP/warmup checks passed. Changed-file lint and formatting pass. Fresh HF A100 smokes completed two updates and checkpoint-2 reload in Harbor, native OpenCode and whitebox. Reload evaluation graded 4/4, 4/4 and 1/1 cases respectively. Harbor produced a nonzero gradient; the other tiny groups had no reward contrast. FineEnvs validation records exact pins and artifacts. Also merged upstream #7460 to fix the development-dependency CI failures from stale Qwen MoE xfail markers.
Before submitting
AI writing disclosure
Note
Medium Risk
Changes the experimental loop-owning GRPO ingestion path and removes turn-selection hooks, so training correctness now depends on OpenEnv capture contracts; mitigated by extensive contract/reconciliation tests and pinned OpenEnv revision.
Overview
Loop-owning AsyncGRPO harness training now consumes OpenEnv’s validated
TrainingTrace(fetch_training_trace()) instead of rebuilding rows from proxy traces and tokenizer re-encoding. OpenEnv’s full-sequenceloss_maskmaps toTurnRecord.output_mask, so partial supervision and rewritten-history forks are preserved in reconciliation (loop-owning mode defaults to lossless capture withfork_threshold_tokens=0).The
HarnessRolloutWorkerAPI dropstrain_turn_fn/agent_turn_fn; turn selection and auxiliary-call filtering move into OpenEnv masks. Factories are built via a callable that receives the trainer’ssampling=policy (e.g.functools.partial(HarborSessionFactory, ...)). Invalid captures raiseCaptureContractErrorand stop the worker; transport failures stay unscorable.Adds the
async_grpo_harborexample (local + HF Jobs launcher), pins OpenEnv86a180e…, and updates OpenCode examples/docs accordingly. Packaged contract tests intests/experimental/test_openenv_tito.pycover mask round-trips, sampling injection, and failure modes.Reviewed by Cursor Bugbot for commit 2ff1235. Bugbot is set up for automated code reviews on this repo. Configure here.